//===- Hello.cpp - Example code from "Writing an LLVM Pass" ---------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements two versions of the LLVM "Hello World" pass described
// in docs/WritingAnLLVMPass.html
//
//===----------------------------------------------------------------------===//

#include <cassert>
#include <iostream>
#include <llvm/Bitcode/BitcodeReader.h>
#include <llvm/Bitcode/BitcodeWriter.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/InstIterator.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/LegacyPassManager.h>
#include <llvm/IRReader/IRReader.h>
#include <llvm/Pass.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/SourceMgr.h>
#include <llvm/Support/ToolOutputFile.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/Transforms/Scalar.h>
#include <llvm/Transforms/Utils.h>
#include <set>

using namespace llvm;
static ManagedStatic<LLVMContext> GlobalContext;
static LLVMContext &getGlobalContext() { return *GlobalContext; }
/* In LLVM 5.0, when  -O0 passed to clang , the functions generated with clang
 * will have optnone attribute which would lead to some transform passes
 * disabled, like mem2reg.
 */
struct EnableFunctionOptPass : public FunctionPass {
  static char ID;
  EnableFunctionOptPass() : FunctionPass(ID) {}
  bool runOnFunction(Function &F) override {
    if (F.hasFnAttribute(Attribute::OptimizeNone)) {
      F.removeFnAttr(Attribute::OptimizeNone);
    }
    return true;
  }
};

class getfuncPtrs {
private:
  // 待分析调用对象的Call指令
  CallInst *callinst;
  // 每一个call inst是一个实例，call_list 由getDirectFuncptr填入
  std::set<Function *> call_list;
  // 维护分析过程中的调用栈，便于找到实参传入点
  std::stack<CallInst *> callstack;

public:
  getfuncPtrs(CallInst *inst) : callinst(inst), call_list({}){};
  ~getfuncPtrs(){};

  void getDirectFuncptr(Function *callee);
  void getArgFuncptr(Argument *argument);
  void getPhiNodeFuncptr(PHINode *phi_node);
  void getCallreturnFuncptr(CallInst *callinst);

  void getDirectCallee(CallInst *callinst);
  void getIndirectCallee(CallInst *callinst);

  std::set<Function *> analysis();
};

void getfuncPtrs::getDirectFuncptr(Function *callee) { this->call_list.insert(callee); };

void getfuncPtrs::getArgFuncptr(Argument *argument) {
  unsigned offset = argument->getArgNo();

  // 如果调用栈里面不为空，就说明这个参数是在指定地方被调用的
  if (!this->callstack.empty()) {
    CallInst *ci = this->callstack.top();
    this->callstack.pop();
    Value *funcptr = ci->getArgOperand(offset);
    if (Function *func = dyn_cast<Function>(funcptr)) {
      this->getDirectFuncptr(func);
    } else if (Argument *argument = dyn_cast<Argument>(funcptr)) {
      this->getArgFuncptr(argument);
    } else if (PHINode *phi_node = dyn_cast<PHINode>(funcptr)) {
      this->getPhiNodeFuncptr(phi_node);
    } else if (CallInst *call_inst = dyn_cast<CallInst>(funcptr)) {
      // the function pointer is a call intruction return value,
      // recursively process
      this->getCallreturnFuncptr(call_inst);
    }
  } else {
    // get its parent caller to get the Actual function pointer parameter
    Function *parent_function = argument->getParent();
    // llvm::errs() << "parent_function: " << parent_function->getName() << "\n";
    for (auto &U : parent_function->uses()) {
      // User 虽然逻辑上必须是函数，但是在代码中这个函数可能是phi的结果，
      // 也可能是call的参数，也可能是call的函数指针返回值
      if (CallInst *call_inst = dyn_cast<CallInst>(U.getUser())) {
        // llvm::errs() << "call-inst: " << *call_inst << "\n";
        Value *funcptr = call_inst->getArgOperand(offset);
        if (Function *func = dyn_cast<Function>(funcptr)) {
          this->getDirectFuncptr(func);
        } else if (Argument *argument = dyn_cast<Argument>(funcptr)) {
          this->getArgFuncptr(argument);
        } else if (PHINode *phi_node = dyn_cast<PHINode>(funcptr)) {
          this->getPhiNodeFuncptr(phi_node);
        } else if (CallInst *call_inst = dyn_cast<CallInst>(funcptr)) {
          // the function pointer is a call intruction return value,
          // recursively process
          this->getCallreturnFuncptr(call_inst);
        }
      }
    }
  }
};

void getfuncPtrs::getPhiNodeFuncptr(PHINode *phi_node) {
  for (Value *incoming : phi_node->incoming_values()) {
    if (Function *func = dyn_cast<Function>(incoming)) {
      this->getDirectFuncptr(func);
    } else if (Argument *argument = dyn_cast<Argument>(incoming)) {
      this->getArgFuncptr(argument);
    } else if (PHINode *phi_node = dyn_cast<PHINode>(incoming)) {
      this->getPhiNodeFuncptr(phi_node);
    } else if (CallInst *call_inst = dyn_cast<CallInst>(incoming)) {
      // the function pointer is a call intruction return value,
      // recursively process
      this->getCallreturnFuncptr(call_inst);
    }
  }
};

void getfuncPtrs::getCallreturnFuncptr(CallInst *callinst) {
  // 首先判断被调用的函数指向谁
  getfuncPtrs getfuncptrs(callinst);
  std::set<Function *> call_list = getfuncptrs.analysis();

  this->callstack.push(callinst);
  for (Function *callee : call_list) {
    callinst->setCalledFunction(callee);
    for (inst_iterator inst = inst_begin(callee); inst != inst_end(callee); inst++) {
      if (ReturnInst *ret = dyn_cast<ReturnInst>(&*inst)) {
        Value *ret_val = ret->getReturnValue();
        if (Function *func = dyn_cast<Function>(ret_val)) {
          this->getDirectFuncptr(func);
        } else if (Argument *argument = dyn_cast<Argument>(ret_val)) {
          this->getArgFuncptr(argument);
        } else if (PHINode *phi_node = dyn_cast<PHINode>(ret_val)) {
          this->getPhiNodeFuncptr(phi_node);
        } else if (CallInst *call_inst = dyn_cast<CallInst>(ret_val)) {
          // the function pointer is a call intruction return value,
          // recursively process
          this->getCallreturnFuncptr(call_inst);
        }
      }
    }
  }
};

void getfuncPtrs::getDirectCallee(CallInst *callinst) {
  Function *callee = callinst->getCalledFunction();
  assert(callee != NULL);
  this->getDirectFuncptr(callee);
};

void getfuncPtrs::getIndirectCallee(CallInst *callinst) {
  Value *funcptr = callinst->getCalledValue();
  if (Argument *argument = dyn_cast<Argument>(funcptr)) {
    this->getArgFuncptr(argument);
  } else if (PHINode *phi_node = dyn_cast<PHINode>(funcptr)) {
    this->getPhiNodeFuncptr(phi_node);
  } else if (CallInst *call_inst = dyn_cast<CallInst>(funcptr)) {
    // the function pointer is a call intruction,recursively process
    this->getCallreturnFuncptr(call_inst);
  }
};

// TODO 优化结果展示逻辑
std::set<Function *> getfuncPtrs::analysis() {
  Function *callee = this->callinst->getCalledFunction();
  if (callee == NULL) {
    this->getIndirectCallee(this->callinst);

  } else if (callee->isIntrinsic()) {
    // do nothing
  } else {
    assert(isa<Function>(callee));
    this->getDirectCallee(this->callinst);
  }
  return this->call_list;
};

std::set<Function *> get_call_list(CallInst *callinst) {
  getfuncPtrs getfuncptrs(callinst);
  return getfuncptrs.analysis();
};

// show the result
void display_call_list(unsigned line, const std::set<Function *> &call_list) {
  errs() << line << " : ";
  auto it = call_list.begin();
  auto ie = call_list.end();
  if (it != ie) {
    errs() << (*it)->getName();
    ++it;
  }
  for (; it != ie; ++it) {
    errs() << ", " << (*it)->getName();
  }
  errs() << "\n";
}

char EnableFunctionOptPass::ID = 0;

///!TODO TO BE COMPLETED BY YOU FOR ASSIGNMENT 2
/// Updated 11/10/2017 by fargo: make all functions
/// processed by mem2reg before this pass.
struct FuncPtrPass : public ModulePass {
  static char ID; // Pass identification, replacement for typeid
  FuncPtrPass() : ModulePass(ID) {}

  // <loc, callee> loc: call instruction, callee: called function
  // std::map<unsigned, std::set<StringRef>> insts_call_lists;
  // std::map<unsigned, std::set<Value *>> _insts_call_lists;

  bool runOnModule(Module &M) override {
    for (Module::iterator func = M.begin(); func != M.end(); func++) {
      for (Function::iterator bb = func->begin(); bb != func->end(); bb++) {
        for (BasicBlock::iterator inst = bb->begin(); inst != bb->end(); inst++) {
          if (CallInst *call = dyn_cast<CallInst>(inst)) {
            Function *callee = call->getCalledFunction();
            if (callee == nullptr || !callee->isIntrinsic()) {
              std::set<Function *> call_list = get_call_list(call);
              unsigned line = call->getDebugLoc().getLine();
              display_call_list(line, call_list);
            }
          }
        }
      }
    }
    // analysis finished，no modification, return false
    return false;
  };
};

char FuncPtrPass::ID = 0;
static RegisterPass<FuncPtrPass> X("funcptrpass", "Print function call instruction");

static cl::opt<std::string> InputFilename(cl::Positional, cl::desc("<filename>.bc"), cl::init(""));

int main(int argc, char **argv) {
  LLVMContext &Context = getGlobalContext();
  SMDiagnostic Err;
  // Parse the command line to read the Inputfilename
  cl::ParseCommandLineOptions(argc, argv, "FuncPtrPass \n My first LLVM too which does not do much.\n");

  // Load the input module
  std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
  if (!M) {
    Err.print(argv[0], errs());
    return 1;
  }

  llvm::legacy::PassManager Passes;

  /// Remove functions' optnone attribute in LLVM5.0
  Passes.add(new EnableFunctionOptPass());
  /// Transform it to SSA
  Passes.add(llvm::createPromoteMemoryToRegisterPass());

  /// Your pass to print Function and Call Instructions
  Passes.add(new FuncPtrPass());
  Passes.run(*M.get());
}
